import pandas as pd


def load_file(file_path, label2id=None):
    df = pd.read_csv(file_path, sep='\t', header=None)
    texts = df.iloc[:, 0].tolist()
    labels = df.iloc[:, 1].tolist()
    if label2id is not None:
        labels = [label2id[str(label)] for label in labels]
    return texts, labels


def parse_labels(label_path):
    id2label, label2id = {}, {}
    with open(label_path, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            label = line.strip()
            id2label[i] = label
            label2id[label] = i
    return id2label, label2id


def generate_labels(data_path):
    id2label, label2id = {}, {}
    df = pd.read_csv(data_path, sep='\t', header=None)
    labels = df.iloc[:, 1].tolist()
    max_label = max(labels)

    for i in range(max_label + 1):
        label = i
        id2label[i] = label
        label2id[label] = i
    return id2label, label2id


if __name__ == '__main__':
    file_path = 'data/dev.txt'
    label_path = 'data/label.txt'
    id2label, label2id = parse_labels(label_path)
    texts, labels = load_file(file_path, label2id=label2id)
    print(texts[:5])
    print(labels[:5])
    print(len(texts))
    print(len(labels))
